# Copyright 2023 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import time

import paddle
import paddleaudio
import numpy as np
from paddleaudio.compliance.kaldi import fbank


THRES = 60000
FEAT_FUNC = lambda waveform, sr: fbank(
        waveform=paddle.to_tensor(waveform), 
        sr=sr, 
        frame_shift=10, 
        frame_length=25, 
        n_mels=80)


def audio_trunc(audio_arr, thres=THRES):
    """
    Args:
        audio_arr: paddle.Tensor
        thres: int
    """
    length = audio_arr.shape[1]
    if length > thres:
        audio_arr = audio_arr[:, :thres]
        return audio_arr
    elif length < thres:
        pad_zero = paddle.zeros((1, THRES), dtype=audio_arr.dtype)
        pad_zero[:, :length] = audio_arr
        return pad_zero
    
    
def txt_to_list(gt_file):
    res_list = []
    with open(gt_file, 'r') as file:
        lines = file.readlines()
        for line in lines:
            line = line.strip()
            for digit in line:
                try:
                    num = int(digit)
                    res_list.append(num)
                except ValueError:
                    print(f"Ignoring non-numeric character: {digit}")
    return res_list


def infer_wav(test_wav, model):
    key_test_load = paddleaudio.load(test_wav)
    key_test_load = (audio_trunc(key_test_load[0]), key_test_load[1]) 
    keyword_feat = FEAT_FUNC(*key_test_load) 
    key_input = keyword_feat.unsqueeze(0).numpy() 

    infer_start_time = time.time()
    output = model.execute([key_input])
    infer_end_time = time.time()

    output0 = output[0]
    result0_array = np.array(output0) 
    keyword_score = np.round(np.max(result0_array).item(), 4)
    
    infer_time_in_milliseconds = (infer_end_time - infer_start_time) * 1000
    print(f"Inference time: {infer_time_in_milliseconds:.2f} milliseconds")
    
    return keyword_score
